package us.luosl.lslt.function.stream;

import us.luosl.lslt.collection.AbstractIterator;
import us.luosl.lslt.collection.GroupedIterator;
import us.luosl.lslt.lang.Tuple;
import us.luosl.lslt.lang.Tuple2;

import java.util.*;
import java.util.function.*;
import java.util.stream.Collector;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public interface StreamWrapper<T> {

    /**
     * 获取被包装的 stream
     * @return 被包装的 stream
     */
    Stream<T> getStream();

    default void forEach(Consumer<? super T> action) {
        getStream().forEach(action);
    }

    default void forEachOrdered(Consumer<? super T> action) {
        getStream().forEachOrdered(action);
    }

    default Object[] toArray() {
        return getStream().toArray();
    }

    default <A> A[] toArray(IntFunction<A[]> generator) {
        return getStream().toArray(generator);
    }

    default T reduce(T identity, BinaryOperator<T> accumulator) {
        return getStream().reduce(identity, accumulator);
    }

    default Optional<T> reduce(BinaryOperator<T> accumulator) {
        return getStream().reduce(accumulator);
    }

    default <U> U reduce(U identity, BiFunction<U, ? super T, U> accumulator, BinaryOperator<U> combiner) {
        return getStream().reduce(identity, accumulator, combiner);
    }

    default <R> R collect(Supplier<R> supplier, BiConsumer<R, ? super T> accumulator, BiConsumer<R, R> combiner) {
        return getStream().collect(supplier, accumulator, combiner);
    }

    default <R, A> R collect(Collector<? super T, A, R> collector) {
        return getStream().collect(collector);
    }

    default Optional<T> min(Comparator<? super T> comparator) {
        return getStream().min(comparator);
    }

    default Optional<T> max(Comparator<? super T> comparator) {
        return getStream().max(comparator);
    }

    default long count() {
        return getStream().count();
    }

    default boolean anyMatch(Predicate<? super T> predicate) {
        return getStream().anyMatch(predicate);
    }

    default boolean allMatch(Predicate<? super T> predicate) {
        return getStream().anyMatch(predicate);
    }

    default boolean noneMatch(Predicate<? super T> predicate) {
        return getStream().noneMatch(predicate);
    }

    default Optional<T> findFirst() {
        return getStream().findFirst();
    }

    default Optional<T> findAny() {
        return getStream().findAny();
    }

    default AbstractIterator<T> iterator() {
        return AbstractIterator.create(getStream().iterator());
    }

    default Spliterator<T> spliterator() {
        return getStream().spliterator();
    }

    default boolean isParallel() {
        return getStream().isParallel();
    }

    default Stream<T> sequential() {
        return getStream().sequential();
    }

    default Stream<T> parallel() {
        return getStream().parallel();
    }

    default Stream<T> unordered() {
        return getStream().unordered();
    }

    default Stream<T> onClose(Runnable closeHandler) {
        return getStream().onClose(closeHandler);
    }

    default void close() {
        getStream().close();
    }

    // ADD API

    default List<T> toList(){
        return getStream().collect(Collectors.toList());
    }

    default Set<T> toSet(){
        return getStream().collect(Collectors.toSet());
    }

    default <C extends Collection<T>> C toCollection(Supplier<C> collectionFactory){
        return getStream().collect(Collectors.toCollection(collectionFactory));
    }

    default <K, V> Map<K, V> toMap(Function<T, K> keyExtractor, Function<T, V> valueExtractor) {
        return toMap(keyExtractor, valueExtractor, (v1, v2) -> v2);
    }

    default <K, V> Map<K, V> toMap(Function<T, K> keyExtractor, Function<T, V> valueExtractor, BinaryOperator<V> mergeFunction) {
        return toMap(keyExtractor, valueExtractor, mergeFunction, HashMap::new);
    }

    default <K, V, M extends Map<K, V>> Map<K, V> toMap(Function<T, K> keyExtractor, Function<T, V> valueExtractor, BinaryOperator<V> mergeFunction, Supplier<M> mapSupplier) {
        return getStream().collect(Collectors.toMap(keyExtractor, valueExtractor, mergeFunction, mapSupplier));
    }

    /**
     * 取 stream 中,第一个满足条件的元素，当 stream 中没有满足条件的元素时，会返回 None
     * 例如：
     * Optional<String> r = StreamWrapper.of("3", "2", "1", "1").takeWhile("2"::equals);
     * System.out.println(r);
     * 输出: Optional[2]
     * @param predicate 判定条件
     * @return 找到的结果
     */
    default Optional<T> takeWhile(Predicate<? super T> predicate){
        return getStream().filter(predicate).findFirst();
    }

    /**
     * 从左开始折叠整个 stream
     * 例如,使用 foldLeft 求和:
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3);
     * Integer r = stream.foldLeft(0, (collector, item) -> collector + item);
     * System.out.println(r);
     * 输出: 6
     * @param zero 起始元素
     * @param f 折叠的操作
     * @param <U> 折叠元素类型
     * @return 在折叠后的结果
     */
    default <U> U foldLeft(U zero, BiFunction<U, T, U> f){
        U xs = zero;
        Iterator<T> itr = iterator();
        while(itr.hasNext()){
            xs = f.apply(xs, itr.next());
        }
        return xs;
    }

    /**
     * 对 stream 进行分组
     * 例如，按对2的求余结果进行分组:
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4);
     * Map<Integer, List<Integer>> r = stream.groupBy(i -> i % 2);
     * System.out.println(r);
     * 输出: {0=[2, 4], 1=[1, 3]}
     * @param keyExtractor 分组 key 抽取器
     * @param <K> 分组 key 类型
     * @return
     */
    default <K> Map<K, List<T>> groupBy(Function<T, K> keyExtractor) {
        return getStream().collect(Collectors.groupingBy(keyExtractor));
    }

    /**
     * 对 stream 进行聚合
     * 例如，分组求均值:
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4);
     * Map<Integer, OptionalDouble> r = stream.agg(i -> i % 2, list -> list.stream().mapToInt(i -> i).average());
     * System.out.println(r);
     * 输出: {0=OptionalDouble[3.0], 1=OptionalDouble[2.0]}
     * @param keyExtractor 分组 key 抽取器
     * @param aggMapper 聚合函数
     * @param <K> 分组 key 类型
     * @param <V> 聚合结果 key 类型
     * @return 聚合后的结果
     */
    default <K, V> Map<K, V> agg(Function<T, K> keyExtractor, Function<List<T>, V> aggMapper){
        return wrapEntry(getStream().collect(Collectors.groupingBy(keyExtractor))).mapValues(aggMapper).toMap();
    }

    /**
     * 对流构造字符串
     * 例如:
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4);
     * String r = stream.mkStr(",", i -> "number:" + i, "[", "]");
     * System.out.println(r);
     * 输出: [number:2,number:1,number:3,number:4]
     * @param delimiter 分隔符
     * @param strFun 元素转字符串函数
     * @param prefix 前缀
     * @param suffix 后缀
     * @return 构造的字符串
     */
    default String mkStr(CharSequence delimiter, Function<T, String> strFun, CharSequence prefix, CharSequence suffix){
        return getStream().map(strFun).collect(Collectors.joining(delimiter, prefix, suffix));
    }

    default String mkStr(CharSequence delimiter, Function<T, String> strFun){
        return mkStr(delimiter, strFun, "", "");
    }

    default String mkStr(CharSequence delimiter){
        return mkStr(delimiter, Objects::toString);
    }

    /**
     * 求 stream 中最大的 topN 元素
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4);
     * Collection<Integer> res = stream.maxTopN(2);
     * System.out.println(res);
     * 输出: [3, 4]
     * @param topN 取多少个 top 元素
     * @param comparator 比较器
     * @return topN 集合
     */
    default Collection<T> maxTopN(int topN, Comparator<T> comparator){
        PriorityQueue<T> heap = new PriorityQueue<>(topN, comparator);
        Iterator<T> itr = iterator();
        while (itr.hasNext()){
            T item = itr.next();
            if(heap.size() < topN){
                heap.add(item);
            }else{
                if(comparator.compare(heap.peek(), item) < 0){
                    heap.poll();
                    heap.add(item);
                }
            }
        }
        return heap;
    }

    /**
     * 求 stream 中最小的 topN 元素
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4);
     * Collection<Integer> res = stream.minTopN(2);
     * System.out.println(res);
     * 输出: [1, 2]
     * @param topN 取多少个 top 元素
     * @param comparator 比较器
     * @return topN 集合
     */
    default Collection<T> minTopN(int topN, Comparator<T> comparator){
        return maxTopN(topN, comparator.reversed());
    }

    default Collection<T> maxTopN(int topN){
        @SuppressWarnings("unchecked")
        Comparator<T> comp = (Comparator<T>) Comparator.naturalOrder();
        return maxTopN(topN, comp);
    }

    default Collection<T> minTopN(int topN){
        @SuppressWarnings("unchecked")
        Comparator<T> comp = (Comparator<T>) Comparator.naturalOrder();
        return minTopN(topN, comp);
    }

    /**
     * 滑动该 stream
     * 例如:
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4);
     * Iterator<List<Integer>> itr = stream.sliding(2);
     * while(itr.hasNext()){
     *     System.out.println(itr.next());
     * }
     * 输出:
     * [2, 1]
     * [1, 3]
     * [3, 4]
     * @param size 窗口大小
     * @param step 补偿
     * @return 滑动后的迭代器
     */
    default AbstractIterator<List<T>> sliding(int size, int step){
        Iterator<T> itr = iterator();
        return new GroupedIterator<>(itr, size, step);
    }

    default AbstractIterator<List<T>> sliding(int size){
        Iterator<T> itr = iterator();
        return new GroupedIterator<>(itr, size, 1);
    }

    /**
     * 根据 size 大小，对 stream 进行分组
     * 例如:
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4);
     * Iterator<List<Integer>> r = stream.grouped(2);
     * while(r.hasNext()){
     *      System.out.println(r.next());
     * }
     * 输出:
     * [2, 1]
     * [3, 4]
     * @param size 分组大小
     * @return 分组后的结果
     */
    default Iterator<List<T>> grouped(int size){
        return sliding (size, size);
    }

    /**
     * 分区，同 {@link StreamWrapper#grouped}
     * @param size 分区大小
     * @return 分区后的结果
     */
    default Iterator<List<T>> partition(int size){
        return grouped(size);
    }

    /*** 工厂方法 ***/

    @SafeVarargs
    static <T> GeneralStreamWrapper<T> of(T ...values){
        return wrap(Stream.of(values));
    }

    static <T> GeneralStreamWrapper<T> wrap(Stream<T> stream){
        return () -> stream;
    }

    static <T> GeneralStreamWrapper<T> wrap(Collection<T> collection){
        return wrap(collection.stream());
    }

    static <T> GeneralStreamWrapper<T> wrap(Iterator<T> itr){
        return wrap(AbstractIterator.create(itr).stream());
    }

    static <T> GeneralStreamWrapper<T> wrap(final T[] array) {
        return () -> Arrays.stream(array);
    }

    static <T1, T2> EntryStreamWrapper<T1, T2> wrapEntry(Stream<Tuple2<T1, T2>> stream){
        return () -> stream;
    }

    static <T1, T2, C extends Collection<Tuple2<T1, T2>>> EntryStreamWrapper<T1, T2> wrapEntry(C collection){
        return wrapEntry(collection.stream());
    }

    static <T1, T2> EntryStreamWrapper<T1, T2> wrapEntry(Map<T1, T2> map){
        return wrapEntry(map.entrySet().stream().map(Tuple::of));
    }

}
